import sqlparse
import re
from tqdm import tqdm
from datetime import datetime

def replace_newline_stream(file_path):
    new_content = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            # 移除行尾换行符后添加转义字符
            cleaned_line = line.rstrip('\n') + '\\n'
            new_content.append(cleaned_line)
    return ''.join(new_content)

root = "F:\\客户资料\\南网-党群\\code\\bj-nsbd_2\\可用部署包\\sql\\表数据"

print(f'{datetime.now().strftime("%Y-%m-%d %H:%M:%S")} - 读取源文件')
sql_content = replace_newline_stream(f'{root}\\data.sql')

print(f'{datetime.now().strftime("%Y-%m-%d %H:%M:%S")} - sqlparse拆分语句')
sql_content = sql_content.split("INSERT INTO")
print(f'{datetime.now().strftime("%Y-%m-%d %H:%M:%S")} - SQL数量：{len(sql_content)}')

# 步骤1：定义正则表达式模式，捕获INSERT语句中的schema.table部分
# 模式解释：r"INSERT INTO\s+(\w+\.\w+)" 
#   - \s+ 匹配一个或多个空格
#   - (\w+\.\w+) 捕获组，匹配如"public.users"的schema表名
pattern = re.compile(r'INSERT\s+INTO\s+(?:"([^"]+)"\.|([^\s.]+)\.)?(?:"([^"]+)"|([^\s(]+))', re.IGNORECASE)  # re.IGNORECASE忽略大小写

# 步骤2：初始化字典分组
schema_groups = {}

# 步骤3：遍历SQL语句并分组
jw_cloud_org_count=0
for stmt in tqdm(sql_content, 
             desc="📊 数据处理", 
             unit="条", 
             bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
    stmt = "INSERT INTO " + stmt.strip()
    if '"CLOUD_ADMIN"."jw_cloud_org"' in stmt:
        jw_cloud_org_count+=1
        continue
    match = pattern.search(stmt)  # 搜索匹配项
    if match:
        schema = match.group(1)  # 提取捕获组内容，例如"public.users"
        # 如果schema不在字典中，创建新列表；否则添加到现有列表
        if schema not in schema_groups:
            schema_groups[schema] = []
        schema_groups[schema].append(stmt)
    else:
        schema = 'other'
        if schema not in schema_groups:
            schema_groups[schema] = []
        schema_groups[schema].append(stmt)

print(f'{datetime.now().strftime("%Y-%m-%d %H:%M:%S")} - jw_cloud_org数量：{jw_cloud_org_count}')

# 步骤4：输出分组结果
for schema, statements in tqdm(schema_groups.items(), 
             desc="📊 数据处理", 
             unit="模式", 
             bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
    print(f'{datetime.now().strftime("%Y-%m-%d %H:%M:%S")} - Schema: {schema}')
    if len(statements) > 0:
        with open(f'{root}\\schemas\\{schema}.sql', 'w', encoding='utf-8') as f:
            for stmt in tqdm(statements, 
             desc="📊 数据处理", 
             unit="条", 
             bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
                f.write(f"{stmt}")

print(f'{datetime.now().strftime("%Y-%m-%d %H:%M:%S")} - 完成拆分')